"""
PDF文件解析器
处理 .pdf 文件，支持文本提取和OCR
"""

import io
import logging
from typing import List
from .base_parser import BaseParser

# 导入PDF处理库
try:
    import PyPDF2
    PYPDF2_AVAILABLE = True
except ImportError:
    PYPDF2_AVAILABLE = False

try:
    import fitz  # PyMuPDF
    PYMUPDF_AVAILABLE = True
except ImportError:
    PYMUPDF_AVAILABLE = False

# 导入OCR库
try:
    from PIL import Image
    import pytesseract
    OCR_AVAILABLE = True
except ImportError:
    OCR_AVAILABLE = False

logger = logging.getLogger(__name__)


class PDFParser(BaseParser):
    """PDF文件解析器"""
    
    def __init__(self):
        super().__init__()
        self.supported_extensions = ['.pdf']
    
    def parse_content(self, content: bytes) -> str:
        """解析PDF文件内容"""
        text_content = []
        
        try:
            # 首先尝试使用PyMuPDF（更好的文本提取）
            if PYMUPDF_AVAILABLE:
                try:
                    text_content = self._parse_with_pymupdf(content)
                    if text_content:
                        logger.info("使用PyMuPDF成功解析PDF")
                except Exception as e:
                    logger.warning(f"PyMuPDF解析失败: {str(e)}")
            
            # 如果PyMuPDF不可用或失败，使用PyPDF2
            if not text_content and PYPDF2_AVAILABLE:
                try:
                    text_content = self._parse_with_pypdf2(content)
                    if text_content:
                        logger.info("使用PyPDF2成功解析PDF")
                except Exception as e:
                    logger.warning(f"PyPDF2解析失败: {str(e)}")
            
            # 提取图片中的文字（如果支持OCR）
            if OCR_AVAILABLE and PYMUPDF_AVAILABLE:
                try:
                    image_texts = self._extract_images_with_ocr(content)
                    if image_texts:
                        text_content.append("\n[图片内容]:")
                        text_content.extend(image_texts)
                        logger.info(f"OCR识别到{len(image_texts)}个图片文本")
                except Exception as e:
                    logger.warning(f"PDF图片OCR识别失败: {str(e)}")
            
            # 合并所有文本内容
            result = "\n".join(text_content)
            
            if not result.strip():
                return "PDF文档解析完成，但未提取到文本内容"
            
            # 清理和提取有意义的内容
            cleaned_result = self.extract_meaningful_content(result)
            
            return cleaned_result if cleaned_result.strip() else "PDF文档功能测试需求"
            
        except Exception as e:
            logger.error(f"解析PDF文件失败: {str(e)}")
            return "PDF文档解析失败"
    
    def _parse_with_pymupdf(self, content: bytes) -> List[str]:
        """使用PyMuPDF解析PDF"""
        text_content = []
        
        pdf_document = fitz.open(stream=content, filetype="pdf")
        
        for page_num in range(len(pdf_document)):
            page = pdf_document.load_page(page_num)
            page_text = page.get_text()
            
            if page_text.strip():
                text_content.append(page_text.strip())
                logger.debug(f"提取第{page_num+1}页文本，长度: {len(page_text)}")
        
        pdf_document.close()
        return text_content
    
    def _parse_with_pypdf2(self, content: bytes) -> List[str]:
        """使用PyPDF2解析PDF"""
        text_content = []
        
        pdf_reader = PyPDF2.PdfReader(io.BytesIO(content))
        
        for page_num, page in enumerate(pdf_reader.pages):
            page_text = page.extract_text()
            
            if page_text.strip():
                text_content.append(page_text.strip())
                logger.debug(f"提取第{page_num+1}页文本，长度: {len(page_text)}")
        
        return text_content
    
    def _extract_images_with_ocr(self, content: bytes) -> List[str]:
        """从PDF中提取图片并进行OCR识别"""
        if not (OCR_AVAILABLE and PYMUPDF_AVAILABLE):
            return []
        
        image_texts = []
        
        try:
            pdf_document = fitz.open(stream=content, filetype="pdf")
            
            for page_num in range(len(pdf_document)):
                page = pdf_document.load_page(page_num)
                image_list = page.get_images()
                
                for img_index, img in enumerate(image_list):
                    try:
                        # 获取图片数据
                        xref = img[0]
                        pix = fitz.Pixmap(pdf_document, xref)
                        
                        if pix.n - pix.alpha < 4:  # 确保是RGB或灰度图
                            # 转换为PIL Image
                            img_data = pix.tobytes("png")
                            pil_image = Image.open(io.BytesIO(img_data))
                            
                            # OCR识别
                            ocr_text = pytesseract.image_to_string(pil_image, lang='chi_sim+eng')
                            
                            if ocr_text.strip():
                                image_texts.append(f"图片{page_num+1}-{img_index+1}: {ocr_text.strip()}")
                                logger.debug(f"OCR识别第{page_num+1}页图片{img_index+1}，文本长度: {len(ocr_text)}")
                        
                        pix = None  # 释放内存
                        
                    except Exception as e:
                        logger.debug(f"OCR识别第{page_num+1}页图片{img_index+1}失败: {str(e)}")
                        continue
            
            pdf_document.close()
            
        except Exception as e:
            logger.warning(f"PDF图片提取失败: {str(e)}")
        
        return image_texts
    
    def get_pdf_info(self, content: bytes) -> dict:
        """获取PDF文档信息"""
        info = {
            "pages": 0,
            "has_text": False,
            "has_images": False,
            "file_size": len(content)
        }
        
        try:
            if PYMUPDF_AVAILABLE:
                pdf_document = fitz.open(stream=content, filetype="pdf")
                info["pages"] = len(pdf_document)
                
                # 检查是否有文本和图片
                for page_num in range(min(3, len(pdf_document))):  # 只检查前3页
                    page = pdf_document.load_page(page_num)
                    
                    if page.get_text().strip():
                        info["has_text"] = True
                    
                    if page.get_images():
                        info["has_images"] = True
                    
                    if info["has_text"] and info["has_images"]:
                        break
                
                pdf_document.close()
                
            elif PYPDF2_AVAILABLE:
                pdf_reader = PyPDF2.PdfReader(io.BytesIO(content))
                info["pages"] = len(pdf_reader.pages)
                
                # 检查前几页是否有文本
                for page in pdf_reader.pages[:3]:
                    if page.extract_text().strip():
                        info["has_text"] = True
                        break
        
        except Exception as e:
            logger.warning(f"获取PDF信息失败: {str(e)}")
        
        return info
